package com.ck.framework.aop;

import com.ck.common.annotation.RequestLimiter;
import com.ck.common.exception.BusinessException;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author ck
 * @date 2023/11/28 16:22
 */


/*** Guava 提供的限流工具类RateLimiter
 * 常用方法：
 * 　　　　create（Double permitsPerSecond）方法根据给定的（令牌:单位时间（1s））比例为令牌生成速率
 * 　　　　tryAcquire（）方法尝试获取一个令牌，立即返回true/false，不阻塞，重载方法具备设置获取令牌个数、获取最大等待时间等参数
 * 　　　　acquire（）方法与tryAcquire类似，但是会阻塞，尝试获取一个令牌，没有时则阻塞直到获取成功
 */
@Aspect
@Component
@Slf4j
public class RequestLimiterAspect {
    /**
     * 不同的方法存放不同的令牌桶
     */
    private final Map<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();

    @Around("@annotation(com.ck.common.annotation.RequestLimiter)")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable{
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        //拿limit的注解
        RequestLimiter limit = method.getAnnotation(RequestLimiter.class);
        if (limit != null) {
            //key作用：不同的接口，不同的流量控制
            String key=limit.key();
            RateLimiter rateLimiter = null;
            //验证缓存是否有命中key
            if (!rateLimiterMap.containsKey(key)) {
                // 创建令牌桶
                rateLimiter = RateLimiter.create(limit.permitsPerSecond());
                rateLimiterMap.put(key, rateLimiter);
                log.info("新建了令牌桶={}，容量={}",key,limit.permitsPerSecond());
            }
            rateLimiter = rateLimiterMap.get(key);
            // 拿令牌
            boolean acquire = rateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
            // 拿不到令牌，直接返回异常提示
            if (!acquire) {
                log.error("令牌桶={}，获取令牌失败",key);
                this.responseFail(limit.msg());
                throw new BusinessException("令牌桶获取令牌失败");
            }
        }
        return joinPoint.proceed();
    }


    /**
     * 直接向前端抛出异常
     * @param msg
     */
    private void responseFail(String msg) throws IOException {
        //这里的实现返回，可以根据自己项目返回对应Json或者对象
        HttpServletResponse response=((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getResponse();
        response.setContentType("text/plain;charset=UTF-8");
        PrintWriter writer = response.getWriter();
        writer.println(msg);
        writer.close();
    }


}
